package com.pethospital.config;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.pethospital.dto.MessageDTO;
import com.pethospital.entity.Message;
import com.pethospital.service.ConsultService;
import com.pethospital.util.JwtUtil;
import io.jsonwebtoken.Claims;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;
import org.springframework.web.socket.server.support.DefaultHandshakeHandler;
import org.springframework.web.util.UriComponentsBuilder;

import javax.annotation.Resource;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

/**
 * WebSocket STOMP配置类
 */
@Slf4j
@Configuration
@EnableWebSocketMessageBroker
public class WebSocketConfig implements WebSocketMessageBrokerConfigurer {

    @Resource
    private JwtUtil jwtUtil;
    
    @Resource
    private ConsultService consultService;
    
    @Resource
    private ObjectMapper objectMapper;
    
    @Value("${server.servlet.context-path:}")
    private String contextPath;

    @Override
    public void configureMessageBroker(MessageBrokerRegistry config) {
        // 设置消息代理的前缀，客户端订阅的地址前缀
        config.enableSimpleBroker("/topic");
        // 设置客户端发送消息的前缀
        config.setApplicationDestinationPrefixes("/app");
        
        log.info("WebSocket消息代理配置完成，订阅前缀: /topic, 发送前缀: /app, 上下文路径: {}", contextPath);
    }

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        // 注册STOMP端点，客户端连接的入口点
        log.info("注册WebSocket端点，上下文路径: {}", contextPath);

        // 移除前导斜杠，防止双重斜杠
        String path = contextPath;
        if (path != null && path.startsWith("/")) {
            path = path.substring(1);
        }
        
        // 输出将要注册的完整端点路径
        String wsEndpoint = "/ws";
        String chatWsEndpoint = "/chat/ws";
        
        log.info("注册WebSocket端点: {}", wsEndpoint);
        log.info("注册WebSocket聊天端点: {}", chatWsEndpoint);
        
        registry.addEndpoint(wsEndpoint)
                .setAllowedOriginPatterns("*") // 允许跨域
                .addInterceptors(new WebSocketHandshakeInterceptor()) // 添加拦截器
                .withSockJS(); // 允许使用SockJS
        
        // 添加一个不使用SockJS的端点，用于直接WebSocket连接
        registry.addEndpoint(chatWsEndpoint)
                .setAllowedOriginPatterns("*")
                .addInterceptors(new WebSocketHandshakeInterceptor());
        
        log.info("WebSocket端点注册完成");
    }
    
    /**
     * WebSocket握手拦截器
     * 用于在握手前处理token认证
     */
    public class WebSocketHandshakeInterceptor implements HandshakeInterceptor {
        
        @Override
        public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, 
                                    WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
            log.info("WebSocket握手请求: {}", request.getURI());
            
            try {
                // 从查询参数中获取token
                String query = request.getURI().getQuery();
                Map<String, String> queryParams = UriComponentsBuilder.newInstance()
                        .query(query)
                        .build()
                        .getQueryParams()
                        .toSingleValueMap();
                
                String token = queryParams.get("token");
                log.info("WebSocket Token: {}", token != null ? (token.length() > 10 ? token.substring(0, 10) + "..." : token) : "null");
                
                if (token == null || token.isEmpty()) {
                    log.warn("WebSocket握手失败: 未提供token");
                    return false;
                }
                
                // 验证token
                Claims claims = jwtUtil.getClaimsByToken(token);
                if (claims == null) {
                    log.warn("WebSocket握手失败: token无效");
                    return false;
                }
                
                Long userId = claims.get("userId", Long.class);
                Integer userType = claims.get("userType", Integer.class);
                
                if (userId == null) {
                    log.warn("Token中未包含userId");
                    return false;
                }
                
                log.info("WebSocket连接认证成功，用户ID: {}, 用户类型: {}", userId, userType);
                attributes.put("userId", userId);
                attributes.put("userType", userType);
                
                // 获取会话ID
                String conversationId = queryParams.get("consultId");
                if (conversationId != null) {
                    attributes.put("conversationId", conversationId);
                    log.info("WebSocket连接会话ID: {}", conversationId);
                } else {
                    log.warn("WebSocket连接缺少consultId参数");
                }
                
                return true;
            } catch (Exception e) {
                log.error("WebSocket握手处理异常", e);
                return false;
            }
        }
        
        @Override
        public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, 
                                 WebSocketHandler wsHandler, Exception exception) {
            // 握手后的处理
            if (exception != null) {
                log.error("WebSocket握手异常", exception);
            } else {
                String uri = request.getURI().toString();
                log.info("WebSocket握手成功, URI: {}", uri);
            }
        }
    }
} 